package skewreduce.framework.planner.search;

import java.util.SortedSet;
import java.util.TreeSet;

import org.apache.hadoop.conf.Configuration;

import skewreduce.framework.planner.PartitionNode;
import skewreduce.framework.planner.PriorityScheduler;
import skewreduce.lib.KahanSum;


public class GreedyGlobalScheduleDistribution extends SearchStrategy {
	PriorityScheduler scheduler;
	SortedSet<PartitionNode> currentPlan = new TreeSet<PartitionNode>(); // for scheduler
	double bestCost;
	int numMachines;
	double threshold;
	
	KahanSum x =  new KahanSum();
	KahanSum x2 = new KahanSum();
	int numLeaves;

	@Override
	public void init(PartitionNode root) {
		currentPlan.add(root);
		bestCost = root.getWorkSecs();
		x.add(bestCost);
		x2.add(bestCost*bestCost);
		++numLeaves;
	}

	@Override
	public void reset() {
		currentPlan.clear();
		bestCost = Double.MAX_VALUE;
		x.reset();
		x2.reset();
		numLeaves = 0;
	}
	
	private double getAvg() {
		return x.value() / currentPlan.size();
	}
	
	private double getStdev() {
		double avg = getAvg();
		return Math.sqrt(x2.value() / currentPlan.size() - avg*avg);
	}
	
	private boolean shouldForcePartition(double c) {
		return getAvg() + getStdev() * threshold < c;
	}
	
	@Override
	public boolean shouldPartition(PartitionNode p, PartitionNode l, PartitionNode r,boolean force) {
		scheduler.reset(currentPlan);
		scheduler.add(l);
		scheduler.add(r);
		double schedCost = scheduler.estimateSchedule();
		
		boolean isGood = schedCost < bestCost || force || shouldForcePartition(p.getWorkSecs());
		
		if ( LOG.isDebugEnabled() ) {
			LOG.debug("best cost = {}, current = {}, force = {}, avg = {}, std = {}, n = {}",new Object[] { bestCost, schedCost, force, getAvg(), getStdev(), numLeaves });
		}
		
		if ( isGood ) {
			bestCost = schedCost;
			if ( LOG.isDebugEnabled() ) {
				LOG.debug("Machine {}: partition {} - total cost = {}",new Object[] { numMachines, p.getPartition().toSpec(), bestCost });
			}
			
			double lr = l.getWorkSecs() + r.getWorkSecs();
			double lrp = lr + p.getWorkSecs();
			double lrP = lr - p.getWorkSecs();
			
			x.add(lrP);
			x2.add(lrp * lrP);
			x2.add(- 2.0*l.getWorkSecs()*r.getWorkSecs());
			++numLeaves;

			// add two split nodes to leaves and insert to the queue
			currentPlan.add(l);
			currentPlan.add(r);
		}
		
		return isGood;
	}

	@Override
	public Configuration getConf() {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public void setConf(Configuration conf) {
		scheduler = new PriorityScheduler(conf);
		threshold = conf.getFloat("skewreduce.planner.search.greedyDist.threshold", 1.0f);
	}

	@Override
	public void setNumMachines(int n) {
		this.numMachines = n;
		scheduler.setNumSlots(n);
	}
}
